#include "optimizer_adapter/main.h"

#include "common/config.h"
#include "common/defer.h"
#include "common/formator.h"
#include "common/memory.h"
#include "nodes/nodes.h"
#include "optimizer_adapter/from_pg_query.h"
#include "optimizer_adapter/to_pg_plan.h"
#include "planner/optimizer.h"

extern "C" {
#include <postgres.h>

#include <commands/explain.h>
#include <jit/jit.h>
#include <nodes/params.h>
#include <nodes/parsenodes.h>
#include <optimizer/planner.h>
#include <optimizer/tlist.h>
#include <utils/guc.h>
}

static planner_hook_type prev_planner_hook = nullptr;
static ExplainOneQuery_hook_type prev_explain_hook = nullptr;

namespace optimizer {

static bool explain_memory = false;

MemoryCounter mem_counters;

auto init = []() {
  if constexpr (PLAN_DEBUG)
    MemoryTracer::Instance().Reset();
};

auto dead = []() {
  mem_counters.Count();
  if constexpr (PLAN_DEBUG) {
    LOG_PRINTLN("Memory: used=" UINT64_FORMAT "Bytes  allocated=" UINT64_FORMAT "Bytes",
                (mem_counters.totalspace_ - mem_counters.freespace_), (mem_counters.totalspace_));

    if (const auto leak = MemoryTracer::Instance().Check(); leak != 0)
      elog(WARNING, "memory leak detected: %ld bytes", leak);
  }
};

static PlannedStmt* PgCascadesOptimize(Query* parse, const char* query_string, int cursor_options,
                                       ParamListInfo bound_params) {
  if (OPT_CONFIG(enable_optimizer)) {
    try {
      auto defer = optimizer::ScopedDefer{init, dead};
      auto* query = copyObject(parse);
      return CascadesOptimizer::Optimize(query, query_string, cursor_options, bound_params);
    } catch (const std::exception& e) {
      elog(WARNING, "pg_cascades Failed to plan query, get error: %s", e.what());
    }
  }
  return standard_planner(parse, query_string, cursor_options, bound_params);
}

static void ExplainOneQuery(Query* query, int cursor_options, IntoClause* into, ExplainState* es,
                            const char* query_string, ParamListInfo params, QueryEnvironment* query_env) {
  if (es->memory)
    explain_memory = true;
  prev_explain_hook(query, cursor_options, into, es, query_string, params, query_env);
  if (OPT_CONFIG(enable_optimizer)) {
    ExplainPropertyText("Optimizer", "pg_cascades", es);

    if (explain_memory) {
      auto mem_usedk_b = (mem_counters.totalspace_ - mem_counters.freespace_);
      auto mem_allocatedk_b = (mem_counters.totalspace_);

      if (es->format == EXPLAIN_FORMAT_TEXT) {
        appendStringInfoSpaces(es->str, 2);

        appendStringInfo(es->str, "Memory: used=" UINT64_FORMAT "Bytes  allocated=" UINT64_FORMAT "Bytes", mem_usedk_b,
                         mem_allocatedk_b);
        appendStringInfoChar(es->str, '\n');
      } else {
        ExplainPropertyUInteger("Memory Used", "Bytes", mem_usedk_b, es);
        ExplainPropertyUInteger("Memory Allocated", "Bytes", mem_allocatedk_b, es);
      }
      explain_memory = false;
    }
  }
}

char* CascadesOptimizer::TransformQuery(Query* query) {
  try {
    auto defer = optimizer::ScopedDefer{init, dead};
    QueryToOperatorTransformer transformer;

    const auto logical_tree = transformer.TransformQuery(query);

    return pstrdup(Format::ToString(logical_tree).data());
  } catch (const std::exception& e) {
    elog(ERROR, "pg_cascades Failed to transform query, get error: %s", e.what());
  }
}

PlannedStmt* CascadesOptimizer::Optimize(Query* query, const char* query_string, int cursor_options,
                                         ParamListInfo bound_params) {
  OptimizerContext ctx;
  QueryToOperatorTransformer transformer;

  PropertySet property_set;

  auto logical_exprs = transformer.TransformQuery(query);

  if (logical_exprs == nullptr)
    throw std::runtime_error("Failed to convert query to operator tree");

  optimizer::Optimizer optimizer(&ctx, 1000);

  if (query->sortClause != nullptr) {
  }

  auto root_id = optimizer.BuildPlanTree(std::move(logical_exprs), &property_set);

  PlanGenerator generator{&ctx};

  PlannedStmt* planned_stmt = makeNode(PlannedStmt);

  auto plan = generator.GeneratePlan(root_id, &property_set);

  apply_tlist_labeling(plan.plan_->targetlist, query->targetList);

  planned_stmt->commandType = query->commandType;
  planned_stmt->queryId = query->queryId;
  planned_stmt->hasReturning = (query->returningList != NIL);
  planned_stmt->hasModifyingCTE = query->hasModifyingCTE;
  planned_stmt->canSetTag = query->canSetTag;
  // planned_stmt->transientPlan = glob->transientPlan;
  // planned_stmt->dependsOnRole = glob->dependsOnRole;
  // planned_stmt->parallelModeNeeded = glob->parallelModeNeeded;
  planned_stmt->planTree = plan.plan_;
  /* temporary solution, need to get the final rtable */
  planned_stmt->rtable = plan.rtable_;
  planned_stmt->permInfos = query->rteperminfos;
  // planned_stmt->resultRelations = glob->resultRelations;
  // planned_stmt->appendRelations = glob->appendRelations;
  // planned_stmt->subplans = glob->subplans;
  // planned_stmt->rewindPlanIDs = glob->rewindPlanIDs;
  // planned_stmt->rowMarks = glob->finalrowmarks;
  planned_stmt->relationOids = plan.relation_oids_;
  // planned_stmt->invalItems = glob->invalItems;
  // planned_stmt->paramExecTypes = glob->paramExecTypes;
  planned_stmt->utilityStmt = query->utilityStmt;
  planned_stmt->stmt_location = query->stmt_location;
  planned_stmt->stmt_len = query->stmt_len;
  planned_stmt->jitFlags = PGJIT_NONE;

  return planned_stmt;
}

void CascadesOptimizer::InitGuc() {
  // clang-format off
  DefineCustomBoolVariable(
    "pg_cascades.enable_pg_cascades",
    "Control constant folding during query optimization.",
    "If set to on, try do constant folding during query optimization.",
    &optimizer::OptConfig::Instance().enable_optimizer,
    false,
    PGC_SUSET,
    0,
    nullptr,
    nullptr,
    nullptr
  );
  // clang-format on

  MarkGUCPrefixReserved("pg_cascades");
}

void CascadesOptimizer::InitHooks() {
  prev_planner_hook = planner_hook;
  planner_hook = optimizer::PgCascadesOptimize;

  prev_explain_hook = (ExplainOneQuery_hook != nullptr) ? ExplainOneQuery_hook : standard_ExplainOneQuery;
  ExplainOneQuery_hook = optimizer::ExplainOneQuery;
}

}  // namespace optimizer